package cn.yzfy.framework.starter.common.support;

import java.util.concurrent.BlockingQueue;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Author: 一朝风月
 * @Date: 2025/8/7 9:39
 * @Code 珍琅Code 研发
 * @Version: 1.0
 * @Description 快速消费线程池
 */
public class EagerThreadPoolExecutor extends ThreadPoolExecutor {

    public EagerThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }


    private final AtomicInteger threadNumber = new AtomicInteger(0);

    /**
     * 获取当前线程池已提交的任务数量
     * @return
     */
    public int getSubmittedTaskCount() {
        return threadNumber.get();
    }

    /**
     * 线程执行任务
     * @param command the task to execute
     */
    @Override
    public void execute(Runnable command) {
        // 每次任务执行前，线程数加1
        threadNumber.incrementAndGet();
        try {
            // 执行任务
            super.execute(command);
        } catch (RejectedExecutionException ex) {
            // 如果任务被拒绝，则将任务重新加入队列
            TaskQueue taskQueue = (TaskQueue) super.getQueue();
            try {
                // 将任务重新加入队列失败则抛出异常
                if (!taskQueue.retryOffer(command, 0, TimeUnit.MILLISECONDS)) {
                    threadNumber.decrementAndGet();
                    throw new RejectedExecutionException("Queue capacity is full.", ex);
                }
            } catch (InterruptedException iex) {
                // 线程中断
                threadNumber.decrementAndGet();
                throw new RejectedExecutionException(iex);
            }
        } catch (Exception ex) {
            threadNumber.decrementAndGet();
            throw ex;
        }
    }


    @Override
    protected void afterExecute(Runnable r, Throwable t) {
            // 每次任务执行完毕，线程数减1
             threadNumber.decrementAndGet();
    }
}
